using System;
using System.Reflection;
using NBody;
using NBody.Analysis;
using NBody.Cosmology;
using OptionParser;

[assembly: AssemblyTitle ("NBodyTrim")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Cut particles from an N-body system outside of a certain radius.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class TrimOptions : InputOutputGetOptions
{
    [Option ("-c", "Adjust for centre of mass")]
    public bool CentreOfMass;
    
    [Option ("-r", "Remove particles outside this radius.")]
    public double CutRadius;
    
    [Option ("-x", "X origin.")]
    public double X;
    
    [Option ("-y", "Y origin.")]
    public double Y;
    
    [Option ("-z", "Z origin.")]
    public double Z;
    
    [Option ("-v", "Exclude particles outside virial radius.")]
    public bool ExcludeVirial;
    
    [Option ("-cosmo", "Cosmology for virial radius calculation.")]
    public string CosmoString;
    
    [Option ("-f", "Input file for potential")]
    public string PotFile;
    
    [Option ("-cylinder", "Extract a cylinder instead of a sphere.")]
    public bool Cylinder;
    
    [Option ("-height", "Remove particles outside this cylindrical height.")]
    public double Height;
    
    [Option ("-id", "Extract halo with id from system")]
    public int Id;
    
    [Option ("-fof", "Extract halo with using tipsy marked array file")]
    public string FofFile = "";
    
    public TrimOptions()
    {
        CutRadius = Double.PositiveInfinity;
        X = 0.0;
        Y = 0.0;
        Z = 0.0;
        CosmoString = "SCDM";
        PotFile = "";
    }
}

public class NBodyTrim
{
    public static void Main(string[] args)
    {
        // Process command-line options
        TrimOptions opts = new TrimOptions();
        opts.ProcessArgs(args, "trim");
                        
        // Read in a N-body system from the given file
        NBodySystem s = NBodySystem.Read(opts.InFile);
        
        // first, if wanted, cut out a halo using tipsy file
        if (opts.FofFile != "")
        {
            DataFile.Table fof = DataFile.Table.Read(opts.FofFile);
            NBodySystem news = new NBodySystem(s.NumParts, 276.5333);
            for (int i = 0; i < s.NumParts; i++)
            {
                if (fof[i, 0] == opts.Id)
                    news.AddParticle(s[i]);
            }
            news.AdjustForCentreOfMass(1e-3);
            news.Write(opts.OutFile);
            Environment.Exit(0);
        }
        
        // if given, read in potential file and assign potentials
        if (opts.PotFile != "")
            Energy.SetPotentialFromData(ref s, DataFile.Table.Read(opts.PotFile));
        //else
        //    Energy.CalculatePotentialShell(ref s, 2.0);
                
        // Right now, the user can specifiy to adjust for c/m or not, 
        // but not the tolerance parameter.  I think this is fine.
        if (opts.CentreOfMass)
            s.AdjustForCentreOfMass(1e-3);
        
        // Translate the system to user-specified origin
        if (opts.X != 0.0 || opts.Y != 0.0 || opts.Z != 0.0)
            s.TranslateTo((double)opts.X, (double)opts.Y, (double)opts.Z);
        
        // if the user wishes to exclude particles outside the virial
        // radius, we need a cosmology to calculate the background density
        // at s.Time.  The user can input the cosmology as a string, but
        // it better corresponde to a specific cosmology defined in the 
        // Cosmology class.
        if (opts.ExcludeVirial)
        {
            Model cosmo = Model.FromString(opts.CosmoString);
            double z = cosmo.Redshift(s.Time);
            // this changes the system!
            opts.CutRadius = VirialRadius.CalculateR200(s, cosmo.Density(z));
        }
        
        if (opts.Cylinder)
        {
            s.ExtractCylinder(opts.CutRadius, opts.Height);
        }
        else if (!Double.IsInfinity(opts.CutRadius))
            s.ExtractSphere(opts.CutRadius);
        
        s.Write(opts.OutFile);
                
        // if requested, also output a potential file
        if (opts.PotFile != "")
        {
            DataFile.Table data = new DataFile.Table(s.NumParts, 1, "Potential");
            for (int i = 0; i < s.NumParts; i++)
                data[i, 0] = ((EnergyParticle)s[i]).Potential;
            data.Print(opts.OutFile + ".pot");
        }
        
    }
}
